"""
https://pytorch.org/vision/stable/datasets.html#torchvision.datasets.ImageFolder
"""
import torch as pt
import torchvision as ptv
from torch.utils.data import DataLoader
import matplotlib.pyplot as plt
import os
import sys
import re
import numpy as np

pt.manual_seed(777)

BATCH_SIZE = 16

data_dir = '../../../../large_data/DL2/_many_files/cifar2_fast/train'


def is_valid_file(path):
    ext = os.path.splitext(path)[1].lower()
    if ext == '.jpg':
        return True
    else:
        return False


train_ds = ptv.datasets.ImageFolder(root=data_dir, is_valid_file=is_valid_file, transform=ptv.transforms.ToTensor())
train_dl = DataLoader(dataset=train_ds, batch_size=BATCH_SIZE, shuffle=True)

plt.figure(figsize=[12, 6])
spr = 4
spc = 8
spn = 0


for bx, by in train_dl:
    for i, bxi in enumerate(bx):
        spn += 1
        if (spn > spr * spc):
            break
        plt.subplot(spr, spc, spn)
        bxi = bxi.transpose(0, 2)
        bxi = bxi.transpose(0, 1)
        plt.imshow(bxi)
        plt.axis('off')
        cls_id = by[i].item()
        plt.title(str(cls_id) + ': ' + train_ds.classes[cls_id])
    if (spn > spr * spc):
        break
